import tree_sitter_python as tspython
import tree_sitter_cuda as tscuda
from tree_sitter import Language, Parser

class BasedParser:
    def __init__(self):
        pass

    def get_chunk(self, text:str):
        raise NotImplementedError("This method should be overridden by subclasses")

class CodeParser(BasedParser):
    def __init__(self):
        super().__init__()

    def get_language(self, lang: str):
        if lang == "python":
            return Language(tspython.language())
        elif lang == "cuda":
            return Language(tscuda.language())
        else:
            raise ValueError(f"Unsupported language: {lang}")
    
    def get_chunk(self, code: str, lang: str, max_chunk_size: int = 200, overlap_size: int = 20) -> list:
        parser = Parser(self.get_language(lang))
        tree = parser.parse(bytes(code, "utf8"))
        root_node = tree.root_node
        chunks = []
        for node in root_node.children:
            # extract function and class definitions
            if node.type in ("function_definition", "class_definition"):
                start = node.start_point[0]
                end = node.end_point[0]
                chunk = "\n".join(code.splitlines()[start:end+1])
                if len(chunk.split()) > max_chunk_size:
                    # If chunk is too large, split it further
                    words = chunk.split()
                    for i in range(0, len(words), max_chunk_size - overlap_size):
                        sub_chunk = " ".join(words[i:i + max_chunk_size])
                        chunks.append(sub_chunk)
                else:
                    chunks.append(chunk)
        return chunks